[DeepCompile] Fix KeyError on frozen parameters in ZeRO-3 - #8214

Under ZeRO-3 with "compile": {"deepcompile": true}, engine.compile() fails with a KeyError if the model has any parameter with requires_grad=False. LoRA/PEFT and partial-freeze runs cannot use DeepCompile at all.

init_z3() looks up a grad partition for every module parameter, but the stage-3 optimizer builds that map only for the parameters it owns. _get_trainable_parameter_groups() drops requires_grad=False params (stage3.py L651), and the map is filled from the resulting fp16_groups (stage3.py L706-720), so the first frozen parameter is simply not there.

Repro

# repro.py, run with: torchrun --nproc_per_node=2 repro.py
import torch
import deepspeed


class Net(torch.nn.Module):

    def __init__(self, dim=128):
        super().__init__()
        self.frozen = torch.nn.Linear(dim, dim)
        self.trainable = torch.nn.Linear(dim, dim)
        self.frozen.requires_grad_(False)

    def forward(self, x):
        return self.trainable(self.frozen(x)).sum()


config = {
    "train_batch_size": 2,
    "train_micro_batch_size_per_gpu": 1,
    "optimizer": {"type": "Adam", "params": {"lr": 1e-4}},
    "zero_optimization": {"stage": 3},
    "bf16": {"enabled": True},
    "compile": {"deepcompile": True},
}

model = Net()
trainable = [p for p in model.parameters() if p.requires_grad]
engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=trainable, config=config)
engine.compile()

x = torch.randn(1, 128, device=engine.device, dtype=torch.bfloat16)
loss = engine(x)
engine.backward(loss)
engine.step()
print("ok")

File "deepspeed/compile/init_z3.py", line 131, in init_z3
    grad_buffer = optimizer._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[p.ds_id]
KeyError: 0

ds_id 0 is the frozen layer's weight. We originally hit this through the HF Trainer path with a LoRA adapter, where the crash happens inside accelerator.prepare() when accelerate calls engine.compile(). Also reproduces on the released 0.19.2 and 0.19.3.

Fix

Skip the lookup for frozen parameters and leave them with the empty buffer that the optimizer-less path (use_opt == False) already uses for every parameter.

Nothing reads that buffer for a frozen parameter. add_gather_and_reduce() skips parameters whose grad node is None (passes/zero3_compile.py L130), and the registered buffer is only consumed by flushReduceBucket through those reduce ops (csrc/compile/z3.cpp L261, L282, L311). set_grad_buffer() a few lines below already applies the same requires_grad guard against the same map. Frozen parameters are still registered with the native handle, since they are partitioned and need gather/release ops in the forward graph.

Test

TestDeepCompile::test_frozen_params runs the existing SimpleFrozenModel for 10 steps on 2 ranks with ZeRO-3 + deepcompile, comparing loss and parameters against a ZeRO-0 eager baseline. 10 steps so the run crosses the WARMUP boundary and the prefetch and selective-gather passes also see the frozen parameters. compare_loss() takes an optional model_cls, so the existing callers are unchanged.

Without the fix the test fails at engine.compile() with the same KeyError.

Notes

ZeRO-1/2 does not hit this, since init_z1() iterates optimizer.bit16_groups, which contains only trainable parameters. I have not checked whether frozen parameters work end to end there, so this change is scoped to ZeRO-3.

grad_partitions.get(p.ds_id, torch.Tensor()) would also cover a trainable parameter that was never passed to the optimizer, but that case is already broken in eager ZeRO-3, so I kept the narrower guard to match set_grad_buffer(). Happy to switch if you prefer the wider one.